import os
import pandas as pd
import numpy as np
from sklearn.preprocessing import LabelEncoder, OneHotEncoder
from sklearn.feature_extraction.text import TfidfVectorizer
from config import Config
from scipy.sparse import save_npz, hstack, csr_matrix


def create_user_features(users, interactions):
    """创建用户特征"""
    # 基本特征
    user_features = users[['user_id']].copy()

    # 编码分类特征
    le = LabelEncoder()
    user_features['age_group_encoded'] = le.fit_transform(users['age_group'])
    user_features['gender_encoded'] = le.fit_transform(users['gender'])
    user_features['preference_encoded'] = le.fit_transform(users['preference'])
    user_features['price_sensitivity_encoded'] = le.fit_transform(users['price_sensitivity'])

    # 交互统计特征
    user_stats = interactions.groupby('user_id').agg(
        total_actions=('action_value', 'count'),
        buy_ratio=('action', lambda x: (x == 'buy').mean()),
        avg_action_value=('action_value', 'mean'),
        last_active=('timestamp', 'max')
    ).reset_index()

    user_features = user_features.merge(user_stats, on='user_id', how='left')

    # 处理缺失值
    user_features.fillna({
        'total_actions': 0,
        'buy_ratio': 0,
        'avg_action_value': 0
    }, inplace=True)

    # 最近活跃天数
    max_date = interactions['timestamp'].max()
    user_features['days_since_active'] = (max_date - user_features['last_active']).dt.days
    user_features.drop('last_active', axis=1, inplace=True)

    # 热门品类偏好
    user_cat_pref = interactions.groupby(['user_id', 'category'])['action_value'].sum().unstack(fill_value=0)
    top_categories = user_cat_pref.sum().nlargest(Config.TOP_CATEGORIES).index
    user_cat_pref = user_cat_pref[top_categories]

    # 合并特征
    user_features = user_features.merge(user_cat_pref, on='user_id', how='left')
    user_features.fillna(0, inplace=True)

    return user_features


def create_item_features(products):
    """创建商品特征"""
    item_features = products[['product_id']].copy()

    # 编码分类特征
    le = LabelEncoder()
    item_features['category_encoded'] = le.fit_transform(products['category'])
    item_features['subcategory_encoded'] = le.fit_transform(products['subcategory'])
    item_features['flavor_encoded'] = le.fit_transform(products['flavor'])

    # 数值特征
    item_features['price'] = products['price']
    item_features['sales'] = products['sales']
    item_features['rating'] = products['rating']
    item_features['is_new'] = products['is_new'].astype(int)

    # TF-IDF 商品名称特征
    tfidf = TfidfVectorizer(max_features=50, stop_words='english')
    name_features = tfidf.fit_transform(products['name'])

    # 热门品类特征
    top_categories = products['category'].value_counts().nlargest(Config.TOP_CATEGORIES).index
    for cat in top_categories:
        item_features[f'cat_{cat}'] = (products['category'] == cat).astype(int)

    # 热门口味特征
    top_flavors = products['flavor'].value_counts().nlargest(Config.TOP_FLAVORS).index
    for flavor in top_flavors:
        item_features[f'flavor_{flavor}'] = (products['flavor'] == flavor).astype(int)

    # 组合所有特征
    dense_features = item_features.drop(['product_id'], axis=1).values
    sparse_features = hstack([csr_matrix(dense_features), name_features])

    return item_features, sparse_features


def create_interaction_matrix(interactions, users, products):
    """创建交互矩阵"""
    # 创建映射
    user_ids = users['user_id'].unique()
    item_ids = products['product_id'].unique()

    user_id_map = {id: i for i, id in enumerate(user_ids)}
    item_id_map = {id: i for i, id in enumerate(item_ids)}

    # 创建矩阵
    interaction_matrix = np.zeros((len(user_ids), len(item_ids)))

    for _, row in interactions.iterrows():
        user_idx = user_id_map.get(row['user_id'])
        item_idx = item_id_map.get(row['product_id'])

        if user_idx is not None and item_idx is not None:
            # 累加加权行为值
            interaction_matrix[user_idx, item_idx] += row['weighted_action']

    return csr_matrix(interaction_matrix), user_id_map, item_id_map


def save_features(user_features, item_features, item_feature_matrix, interaction_matrix, user_id_map, item_id_map):
    """保存特征"""
    os.makedirs(Config.PROCESSED_DATA_DIR, exist_ok=True)

    user_features.to_pickle(Config.USER_FEATURES_FILE)
    item_features.to_pickle(Config.ITEM_FEATURES_FILE)
    save_npz(Config.INTERACTION_MATRIX_FILE, interaction_matrix)

    # 保存映射
    pd.Series(user_id_map).to_pickle(os.path.join(Config.PROCESSED_DATA_DIR, "user_id_map.pkl"))
    pd.Series(item_id_map).to_pickle(os.path.join(Config.PROCESSED_DATA_DIR, "item_id_map.pkl"))

    print("Features saved successfully")


def feature_engineering():
    """特征工程流程"""
    # 加载数据
    data = pd.read_pickle(os.path.join(Config.PROCESSED_DATA_DIR, "full_data.pkl"))
    users = pd.read_csv(Config.USERS_FILE)
    products = pd.read_csv(Config.PRODUCTS_FILE)

    # 创建特征
    print("Creating user features...")
    user_features = create_user_features(users, data)

    print("Creating item features...")
    item_features, item_feature_matrix = create_item_features(products)

    print("Creating interaction matrix...")
    interaction_matrix, user_id_map, item_id_map = create_interaction_matrix(
        data, user_features, item_features
    )

    # 保存特征
    save_features(user_features, item_features, item_feature_matrix, interaction_matrix, user_id_map, item_id_map)

    return user_features, item_features, interaction_matrix